-- YUVField

{-
  Convert YUV 4:2:0 image file between frame format (progressive) and field format (interlaced).
  
  Refer to :
    YUV 420 (I420/IYUV) file format : http://www.fourcc.org/yuv.php#IYUV
-}

import Char
import Control.Monad
import qualified Data.ByteString as B
import Data.Word
import Foreign ( unsafePerformIO )
import Text.Regex
import Text.ParserCombinators.Parsec
import List
import System
import System.IO
import qualified System.IO.Error as E


showCmdUsage = concat $ intersperse "\n" [
  "To convert YUV 420 sequence between progressive and interlaced.",
  "",
  "Usage:",
  "\tYUVField.exe command size in-file out-file(s)",
  "",
  "Command:",
  "  Operations from one file to another file :",
  "\tfld2frm                 Convert fields to frames.",
  "\tfrm2fld                 Convert frames to fields.",
  "  Operations between one frame file and two field files :",
  "\tmerge, m                Merge fields to frames.",
  "\tsplit, s                Split frames to fields.",
  "",
  "Size:",
  "  Indicates the image size of source YUV 4:2:0 file.",
  "\t<width>x<height>        Input image size.",
  "\t<mbwidth>m<mbheight>    Input image size in MB(s).",
  "\tcif, qcif, d1           Name for image size.",
  "",
  "File:",
  "  When merging fields, I read file1 and file2, write to file3.",
  "  When spliting frames, read file1 and write into file2 and file3.",
  "",
  "Example:",
  "  YUVField.exe frm2fld d1 mobile_d1.yuv mobile_d1_fld.yuv",
  "  YUVField.exe split 720x576 mobile_d1.yuv fld_top.yuv fld_btm.yuv",
  "",
  "Version 1.0",
  "Made by Levi.G. 2009-03. All rights reserved.",
  "Email: leaveye <dot> guo <at> gmail <dot> com" ]


data CLICommand = CLI_Usage | CLI_Merge | CLI_Split
                | CLI_ConvMerge | CLI_ConvSplit
  deriving ( Enum, Show )


parseArgs args = case args of
    sCmd:sSize:fs@(_:_:_) -> do
            let Right opt = parseCmd sCmd
            let Right (w,h) = parseSize sSize
            if (checkPicSize w && checkPicSize h)
              then return (opt,(w,h),fs)
              else returnUsage
    otherwise   -> returnUsage
  where
    checkPicSize :: Int -> Bool
    checkPicSize n = n `rem` 16 == 0
    returnUsage = return (CLI_Usage,(0,0),[])
    
    parseSize arg = parse parsers "image size" arg
      where
        parsers = sizeInPix <|> sizeCif <|> sizeQcif <|> sizeD1
        sizeInPix :: Parser (Int,Int)
        sizeInPix = do {
            w <- many1 digit;
            f <- do { char 'x'; return 1 } <|>
                 do { char 'm'; return 16 };
            h <- many1 digit;
            return ( read w * f, read h * f ) }
          <?> "\"<width>x<height>\", \"<mbwidth>m<mbheight>\""
        sizeCif :: Parser (Int,Int)
        sizeCif = do string "cif"; return ( 352, 288 )
        sizeQcif :: Parser (Int,Int)
        sizeQcif = do string "qcif"; return ( 176, 144 )
        sizeD1 :: Parser (Int,Int)
        sizeD1 = do try ( string "d1" ) <|> string "done"; return ( 720, 576 )
    
    parseCmd arg = parse parsers "operation" arg
      where
        parsers = pMerge <|> pSplit <|> pF2F <|> pFailed
        pMerge = do { char 'm'; return CLI_Merge } <?> "\"merge\""
        pSplit = do { char 's'; return CLI_Split } <?> "\"split\""
        pF2F = do { try (string "fl" <?> "\"fld2frm\""); return CLI_ConvMerge }
           <|> do { try (string "fr" <?> "\"frm2fld\""); return CLI_ConvSplit }
        pFailed = return CLI_Usage


main = getArgs >>= \args -> do
    (opt,sz,fs) <- parseArgs args
    case opt of
      CLI_Usage     -> putStr showCmdUsage
      CLI_Merge     -> withBinaryFile3
                         (fs!!0) ReadMode (fs!!1) ReadMode
                         (fs!!2) WriteMode
                         (cmdMerge sz)
      CLI_Split     -> withBinaryFile3
                         (fs!!0) ReadMode
                         (fs!!1) WriteMode (fs!!2) WriteMode
                         (cmdSplit sz)
      CLI_ConvMerge -> withBinaryFile2
                         (fs!!0) ReadMode
                         (fs!!1) WriteMode
                         (cmdConvMerge sz)
      CLI_ConvSplit -> withBinaryFile2
                         (fs!!0) ReadMode
                         (fs!!1) WriteMode
                         (cmdConvSplit sz)
  where
    withBinaryFile2 f1 m1 f2 m2 todo =
      withBinaryFile f2 m2 $ \h2 ->
      withBinaryFile f1 m1 $ \h1 ->
      todo h1 h2
    withBinaryFile3 f1 m1 f2 m2 f3 m3 todo =
      withBinaryFile f3 m3 $ \h3 ->
      withBinaryFile f2 m2 $ \h2 ->
      withBinaryFile f1 m1 $ \h1 ->
      todo h1 h2 h3


fetch l h = do
  raw <- B.hGet h l
  return $ raw `B.append` B.replicate (l - B.length raw) 0x80

write = B.hPut

pump :: Int -> Handle -> Handle -> IO ()
pump l from to = fetch l from >>= write to


eofHandler e = if E.isEOFError e then return () else E.ioError e


doUntil :: (IO Bool) -> (IO a) -> (IO b) -> IO b
doUntil cond work end =
    cond >>= \c ->
    if c then end
      else do { work; doUntil cond work end }

anyEof :: [Handle] -> IO Bool
anyEof = liftM or . mapM hIsEOF


cmdMerge :: (Int,Int) -> Handle -> Handle -> Handle -> IO ()
cmdMerge sz top btm frm = doUntil
    (anyEof [top,btm])
    (oneMerge sz top btm frm)
    (hTell frm >>= \p -> putStrLn $ "Written " ++ (show p) ++ " bytes into file.")
  where
    oneMerge :: (Int,Int) -> Handle -> Handle -> Handle -> IO ()
    oneMerge (w,h) f0 f1 t = sequence_ $ zipWith3 pump ls fs ts
      where ls = concatMap (replicate (h+h)) [w,w `div` 2]
            fs = cycle [f0,f1]
            ts = repeat t

cmdSplit :: (Int,Int) -> Handle -> Handle -> Handle -> IO ()
cmdSplit sz frm top btm = doUntil
    (hIsEOF frm)
    (oneSplit sz frm top btm)
    (hTell btm >>= \p -> putStrLn $ "Written " ++ (show p) ++ " bytes into each file.")
  where
    oneSplit :: (Int,Int) -> Handle -> Handle -> Handle -> IO ()
    oneSplit (w,h) f t0 t1 = sequence_ $ zipWith3 pump ls fs ts
      where ls = concatMap (replicate h) [w,w `div` 2]
            fs = repeat f
            ts = cycle [t0,t1]

cmdConvMerge :: (Int,Int) -> Handle -> Handle -> IO ()
cmdConvMerge sz fld frm = doUntil
    (hIsEOF fld)
    (oneConvMerge sz fld frm)
    (hTell frm >>= \p -> putStrLn $ "Written " ++ (show p) ++ " bytes into file.")
  where
    oneConvMerge :: (Int,Int) -> Handle -> Handle -> IO ()
    oneConvMerge (w,h) f t = do
        top <- fetch szPic f
        btm <- fetch szPic f
        mapM_ (write t) [ sub o l d | (o,l) <- zip os ls, d <- [top,btm] ]
      where
        szPic = h * (w + w `div` 2)
        ls = concatMap (replicate h) [w,w `div` 2]
        os = init $ scanl (+) 0 ls
        sub o l = B.take l . B.drop o

cmdConvSplit :: (Int,Int) -> Handle -> Handle -> IO ()
cmdConvSplit sz frm fld = doUntil
    (hIsEOF frm)
    (oneConvSplit sz frm fld)
    (hTell fld >>= \p -> putStrLn $ "Written " ++ (show p) ++ " bytes into each file.")
  where
    oneConvSplit :: (Int,Int) -> Handle -> Handle -> IO ()
    oneConvSplit (w,h) f t = do
        frm <- fetch szPic f
        mapM_ (write t) [ sub o l frm | (o,l) <- zip os ls ]
      where
        szPic = h * (w + w `div` 2)
        ls0 = concatMap (replicate (h `div` 2)) [w,w `div` 2]
        ls = ls0 ++ ls0
        os0 = map (*2) . init $ scanl (+) 0 ls0
        os = os0 ++ zipWith (+) ls0 os0
        sub o l = B.take l . B.drop o
